
import torch
import numpy as np

import torch.nn as nn
import torch.nn.functional as F



LR_A = 0.001    # learning rate for actor
LR_C = 0.002    # learning rate for critic
GAMMA = 0.9     # reward discount
TAU = 0.01      # soft replacement



g_总缓存数量 = 100000
g_训练批次 = 1000


class ANet(nn.Module):   # ae(s)=a
	def __init__(self,s_dim,a_dim):
		super(ANet,self).__init__()
		self.fc1 = nn.Linear(s_dim,30)
		self.fc1.weight.data.normal_(0,0.1) # initialization
		self.out = nn.Linear(30,a_dim)
		self.out.weight.data.normal_(0,0.1) # initialization
	def forward(self,x):
		x = self.fc1(x)
		x = F.relu(x)
		x = self.out(x)
		x = F.tanh(x)
		actions_value = x*2
		return actions_value

class CNet(nn.Module):   # ae(s)=a
	def __init__(self,s_dim,a_dim):
		super(CNet,self).__init__()
		self.fcs = nn.Linear(s_dim,30)
		self.fcs.weight.data.normal_(0,0.1) # initialization
		self.fca = nn.Linear(a_dim,30)
		self.fca.weight.data.normal_(0,0.1) # initialization
		self.out = nn.Linear(30,1)
		self.out.weight.data.normal_(0, 0.1)  # initialization
	def forward(self,s,a):
		x = self.fcs(s)
		y = self.fca(a)
		net = F.relu(x+y)
		actions_value = self.out(net)
		return actions_value





class DDPG(object):
	def __init__(self, a_dim, s_dim, a_bound,):
		self.a_dim, self.s_dim, self.a_bound = a_dim, s_dim, a_bound,
		self.memory = np.zeros((g_总缓存数量, s_dim * 2 + a_dim + 1), dtype=np.float32)
		self.pointer = 0
		#self.sess = tf.Session()
		self.Actor_eval = ANet(s_dim,a_dim)
		self.Actor_target = ANet(s_dim,a_dim)
		self.Critic_eval = CNet(s_dim,a_dim)
		self.Critic_target = CNet(s_dim,a_dim)
		self.ctrain = torch.optim.Adam(self.Critic_eval.parameters(),lr=LR_C)
		self.atrain = torch.optim.Adam(self.Actor_eval.parameters(),lr=LR_A)
		self.loss_td = nn.MSELoss()

	def choose_action(self, s):
		s = torch.unsqueeze(torch.FloatTensor(s), 0)
		return self.Actor_eval(s)[0].detach() # ae（s）

	def learn(self):

		for x in self.Actor_target.state_dict().keys():
			eval('self.Actor_target.' + x + '.data.mul_((1-TAU))')
			eval('self.Actor_target.' + x + '.data.add_(TAU*self.Actor_eval.' + x + '.data)')
		for x in self.Critic_target.state_dict().keys():
			eval('self.Critic_target.' + x + '.data.mul_((1-TAU))')
			eval('self.Critic_target.' + x + '.data.add_(TAU*self.Critic_eval.' + x + '.data)')

		# soft target replacement
		#self.sess.run(self.soft_replace)  # 用ae、ce更新at，ct

		indices = np.random.choice(g_总缓存数量, size=g_训练批次)
		bt = self.memory[indices, :]
		bs = torch.FloatTensor(bt[:, :self.s_dim])
		ba = torch.FloatTensor(bt[:, self.s_dim: self.s_dim + self.a_dim])
		br = torch.FloatTensor(bt[:, -self.s_dim - 1: -self.s_dim])
		bs_ = torch.FloatTensor(bt[:, -self.s_dim:])

		a = self.Actor_eval(bs)
		q = self.Critic_eval(bs,a)  # loss=-q=-ce（s,ae（s））更新ae   ae（s）=a   ae（s_）=a_
		# 如果 a是一个正确的行为的话，那么它的Q应该更贴近0
		loss_a = -torch.mean(q) 
		#print(q)
		#print(loss_a)
		self.atrain.zero_grad()
		loss_a.backward()
		self.atrain.step()

		a_ = self.Actor_target(bs_)  # 这个网络不及时更新参数, 用于预测 Critic 的 Q_target 中的 action
		q_ = self.Critic_target(bs_,a_)  # 这个网络不及时更新参数, 用于给出 Actor 更新参数时的 Gradient ascent 强度
		q_target = br+GAMMA*q_  # q_target = 负的
		#print(q_target)
		q_v = self.Critic_eval(bs,ba)
		#print(q_v)
		td_error = self.loss_td(q_target,q_v)
		# td_error=R + GAMMA * ct（bs_,at(bs_)）-ce(s,ba) 更新ce ,但这个ae(s)是记忆中的ba，让ce得出的Q靠近Q_target,让评价更准确
		#print(td_error)
		self.ctrain.zero_grad()
		td_error.backward()
		self.ctrain.step()

	def store_transition(self, s, a, r, s_):
		transition = np.hstack((s, a, [r], s_))
		index = self.pointer % g_总缓存数量  # replace the old memory with new memory
		self.memory[index, :] = transition
		self.pointer += 1




g_Net = None

def f_创建网络模型(输入维度, 输出维度, 训练批次):
	g_Net = DDPG(输出维度, 输入维度, 训练批次)
	pass


def f_预测动作(s):
	a = g_Net.choose_action(s);
	a = np.clip(np.random.normal(a, var), -1.0, 1.0)
	return list(s);

var = 3  # control exploration

def f_训练模型(s, a, r, s_):
	g_Net.store_transition(s, a, r, s_)

	if g_Net.pointer > g_总缓存数量:
		var *= 0.9999
		g_Net.learn()
	









